package com.chatmcp.mcprouter.model;

import com.chatmcp.mcprouter.dto.ProxyInfoDto;
import com.chatmcp.mcprouter.dto.ServerConfigDto;
import com.chatmcp.mcprouter.service.McpClientService;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * SSE Session for managing client connections
 * Mirrors the Go implementation's SSESession
 */
@Slf4j
@Getter
public class SseSession {
    private final SseEmitter emitter;
    private final BlockingQueue<String> messages;
    private final ServerConfigDto serverConfig;
    private final AtomicBoolean closed;
    
    @Setter
    private ProxyInfoDto proxyInfo;
    
    @Setter
    private McpClientService client;

    /**
     * Create a new SSE session
     * 
     * @param emitter SSE emitter
     * @param serverConfig Server configuration
     * @param proxyInfo Proxy information
     */
    public SseSession(SseEmitter emitter, ServerConfigDto serverConfig, ProxyInfoDto proxyInfo) {
        this.emitter = emitter;
        this.messages = new LinkedBlockingQueue<>(100); // Same queue size as Go version
        this.serverConfig = serverConfig;
        this.proxyInfo = proxyInfo;
        this.closed = new AtomicBoolean(false);
        
        // Set up completion callbacks
        this.emitter.onCompletion(() -> close());
        this.emitter.onTimeout(() -> close());
        this.emitter.onError(e -> {
            log.error("Error in SSE emitter: {}", e.getMessage());
            close();
        });
    }

    /**
     * Get the session key
     * 
     * @return Server key
     */
    public String getKey() {
        return proxyInfo.getServerKey();
    }

    /**
     * Get the server command
     * 
     * @return Server command
     */
    public String getCommand() {
        return proxyInfo.getServerCommand();
    }

    /**
     * Get the session ID
     * 
     * @return Session ID
     */
    public String getSessionId() {
        return proxyInfo.getSessionId();
    }

    /**
     * Send a message to the session
     * 
     * @param message Message to send
     * @return True if successful, false otherwise
     */
    public boolean sendMessage(String message) {
        if (closed.get()) {
            log.debug("Session is closed, cannot send message");
            return false;
        }
        
        try {
            return messages.offer(message);
        } catch (Exception e) {
            log.error("Failed to send message to session: {}", e.getMessage());
            return false;
        }
    }

    /**
     * Close the session and clean up resources
     */
    public void close() {
        if (closed.compareAndSet(false, true)) {
            log.debug("Closing SSE session for {}", getSessionId());
            closeClient();
            
            try {
                emitter.complete();
            } catch (Exception e) {
                log.warn("Error completing SSE emitter: {}", e.getMessage());
            }
        }
    }

    /**
     * Close the client connection
     */
    public void closeClient() {
        if (client != null) {
            try {
                client.disconnect();
                log.debug("Client closed for session {}", getSessionId());
            } catch (Exception e) {
                log.error("Error closing client: {}", e.getMessage());
            }
        }
    }

    /**
     * Check if the session is closed
     * 
     * @return True if closed, false otherwise
     */
    public boolean isClosed() {
        return closed.get();
    }
} 